#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <dirent.h>
#include "utils.h"
#include "config.h"
#include "my_numa.h"

// for NUMA
int _num_numa_nodes = 0;
// mapping NUMA nodes to cpus
int _cpu_to_numa_map[MAX_CPUS];
int _enable_internode_migration = 0;

//extern int numa_bitmask_isbitset(const struct bitmask *bmp, unsigned int n);

void init_numa(int enable_internode_migration) {
	_enable_internode_migration = enable_internode_migration;
	for (int i = 0; i < MAX_CPUS; i++) _cpu_to_numa_map[i] = NOT_SET;

	DIR* node_dir = opendir("/sys/devices/system/node");
	if (node_dir == NULL) {
		DBG(DNUMA, "%s", "NUMA support is not available");
		return;
	}
	struct dirent* node_dirent;
	while ((node_dirent = readdir(node_dir)) != NULL) {
		if (strncmp(node_dirent->d_name, "node", 4) == 0) {
			_num_numa_nodes++;
			int node_num = atoi(node_dirent->d_name + 4);
			char dir_name[512] = "/sys/devices/system/node/";
			strcat(dir_name, node_dirent->d_name);
			DIR* cpu_dir = opendir(dir_name);
			if (cpu_dir == NULL) FAIL("Could not open directory for %s\n", dir_name);
			struct dirent* cpu_dirent;
			char cpu_list[255] = "";
			while ((cpu_dirent = readdir(cpu_dir)) != NULL) {
				if (strncmp(cpu_dirent->d_name, "cpu", 3) == 0) {
					if (cpu_dirent->d_name[3] < '0' || cpu_dirent->d_name[3] > '9') continue;
					int cpu_num = atoi(cpu_dirent->d_name + 3);
					strcat(cpu_list, cpu_dirent->d_name + 3);
					strcat(cpu_list, " ");
					_cpu_to_numa_map[cpu_num] = node_num;
				}
			}
			CHECK_CALL(closedir(cpu_dir));
			DBG(DNUMA, "Node %d, cpus %s", node_num, cpu_list);
		}
	}
	CHECK_CALL(closedir(node_dir));
	DBG(DNUMA, "Number of NUMA nodes %d", _num_numa_nodes);
}

int get_num_numa_nodes() {
	if (_enable_internode_migration) return 1;
	return _num_numa_nodes;
}

int same_numa_node(int cid1, int cid2) {
	if (_enable_internode_migration) return 1;
	if (_cpu_to_numa_map[cid1] == _cpu_to_numa_map[cid2]) return 1;
	return 0;
}

int get_numa_node(int cid) {
	if (_enable_internode_migration) return 0;
	return _cpu_to_numa_map[cid];
}
